import subprocess
import sys
import os
import time
from datetime import datetime

# 颜色代码
class Colors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKCYAN = '\033[96m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'

def log(message: str, level: str = "INFO"):
    """打印带时间戳和颜色的日志"""
    timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    
    # 根据日志级别选择颜色
    color_map = {
        "INFO": Colors.OKBLUE,
        "SUCCESS": Colors.OKGREEN,
        "WARNING": Colors.WARNING,
        "ERROR": Colors.FAIL,
        "HEADER": Colors.HEADER
    }
    
    color = color_map.get(level, Colors.OKBLUE)
    print(f'{color}[{timestamp}] [{level}] {message}{Colors.ENDC}')

def print_banner():
    """打印停止横幅"""
    banner = f"""
{Colors.HEADER}{'='*60}
    煤炉监控系统 - 服务停止脚本
    停止所有相关的Python进程
    Version: 1.0.0
{'='*60}{Colors.ENDC}
"""
    print(banner)

def get_python_processes_detailed():
    """获取详细的Python进程信息（包含命令行）"""
    try:
        if os.name == 'nt':  # Windows
            # 使用PowerShell获取进程信息
            ps_script = """
            Get-WmiObject Win32_Process | Where-Object {$_.Name -eq "python.exe"} |
            Select-Object ProcessId, CommandLine |
            ConvertTo-Json
            """

            result = subprocess.run([
                'powershell', '-Command', ps_script
            ], capture_output=True, text=True, check=True)

            processes = []
            if result.stdout.strip():
                import json
                try:
                    data = json.loads(result.stdout)
                    if isinstance(data, dict):
                        data = [data]  # 单个进程转为列表

                    for proc in data:
                        if proc and 'ProcessId' in proc:
                            processes.append({
                                'Id': proc['ProcessId'],
                                'CommandLine': proc.get('CommandLine', ''),
                                'ProcessName': 'python'
                            })
                except json.JSONDecodeError:
                    log("解析进程信息失败", "ERROR")

            return processes
        else:  # Unix/Linux
            result = subprocess.run(['ps', 'axo', 'pid,command'], capture_output=True, text=True, check=True)
            processes = []
            for line in result.stdout.split('\n')[1:]:
                if 'python' in line:
                    parts = line.strip().split(None, 1)
                    if len(parts) >= 2:
                        try:
                            pid = int(parts[0])
                            cmdline = parts[1]
                            processes.append({
                                'Id': pid,
                                'CommandLine': cmdline,
                                'ProcessName': 'python'
                            })
                        except ValueError:
                            continue
            return processes
    except Exception as e:
        log(f"获取进程列表失败: {e}", "ERROR")
        return []

def categorize_processes(processes):
    """将进程按类型分类"""
    categories = {
        'fastapi': [],
        'celery_worker': [],
        'celery_beat': [],
        'start_script': [],
        'other': []
    }

    for proc in processes:
        cmdline = proc.get('CommandLine', '').lower()
        pid = proc.get('Id')

        if 'uvicorn' in cmdline and 'main:app' in cmdline:
            categories['fastapi'].append(proc)
        elif 'celery' in cmdline and 'worker' in cmdline:
            categories['celery_worker'].append(proc)
        elif 'celery' in cmdline and 'beat' in cmdline:
            categories['celery_beat'].append(proc)
        elif 'start_all.py' in cmdline:
            categories['start_script'].append(proc)
        else:
            categories['other'].append(proc)

    return categories

def display_processes_by_category(categories):
    """按分类显示进程"""
    total_count = 0

    for category, processes in categories.items():
        if processes:
            category_names = {
                'fastapi': 'FastAPI 服务',
                'celery_worker': 'Celery Worker 进程',
                'celery_beat': 'Celery Beat 调度器',
                'start_script': '启动脚本',
                'other': '其他Python进程'
            }

            log(f"\n{category_names.get(category, category)}:", "HEADER")
            for proc in processes:
                pid = proc.get('Id')
                cmdline = proc.get('CommandLine', '')
                # 截断过长的命令行
                if len(cmdline) > 80:
                    cmdline = cmdline[:77] + "..."
                log(f"  PID: {pid} - {cmdline}", "INFO")
                total_count += 1

    return total_count

def stop_processes_by_category(categories, selected_categories):
    """按分类停止进程"""
    all_pids = []
    stopped_services = []

    for category in selected_categories:
        if category in categories and categories[category]:
            pids = [proc['Id'] for proc in categories[category]]
            all_pids.extend(pids)

            category_names = {
                'fastapi': 'FastAPI 服务',
                'celery_worker': 'Celery Worker 进程',
                'celery_beat': 'Celery Beat 调度器',
                'start_script': '启动脚本',
                'other': '其他Python进程'
            }
            stopped_services.append(category_names.get(category, category))
            log(f"准备停止 {category_names.get(category, category)}: {pids}", "INFO")

    if not all_pids:
        log("没有找到需要停止的进程", "INFO")
        return False

    try:
        if os.name == 'nt':  # Windows
            cmd = ['powershell', '-Command', f'Stop-Process -Id {",".join(map(str, all_pids))} -Force']
        else:  # Unix/Linux
            cmd = ['kill', '-TERM'] + [str(pid) for pid in all_pids]

        result = subprocess.run(cmd, capture_output=True, text=True)

        if result.returncode == 0:
            log(f"成功停止 {len(all_pids)} 个进程", "SUCCESS")
            for service in stopped_services:
                log(f"  ✓ {service}", "SUCCESS")
            return True
        else:
            log(f"停止进程时出现问题: {result.stderr}", "WARNING")
            return False

    except Exception as e:
        log(f"停止进程失败: {e}", "ERROR")
        return False

def show_menu():
    """显示操作菜单"""
    menu = """
请选择要停止的服务 (可多选，用逗号分隔):
1. FastAPI 服务
2. Celery Worker 进程 (监控任务处理)
3. Celery Beat 调度器 (定时任务调度)
4. 启动脚本
5. 其他Python进程
6. 停止所有Celery监控系统 (Worker + Beat)
7. 停止完整系统 (FastAPI + Celery)
8. 停止所有Python进程
0. 退出

输入选项 (例如: 1,2 或 6): """

    return input(menu).strip()

def parse_menu_selection(selection, categories):
    """解析菜单选择"""
    selected_categories = []

    try:
        choices = [choice.strip() for choice in selection.split(',')]

        for choice in choices:
            if choice == '0':
                return None  # 退出
            elif choice == '1':
                selected_categories.append('fastapi')
            elif choice == '2':
                selected_categories.append('celery_worker')
            elif choice == '3':
                selected_categories.append('celery_beat')
            elif choice == '4':
                selected_categories.append('start_script')
            elif choice == '5':
                selected_categories.append('other')
            elif choice == '6':  # 停止所有Celery监控系统
                selected_categories.extend(['celery_worker', 'celery_beat'])
            elif choice == '7':  # 停止完整系统
                selected_categories.extend(['fastapi', 'celery_worker', 'celery_beat'])
            elif choice == '8':  # 停止所有Python进程
                selected_categories.extend(['fastapi', 'celery_worker', 'celery_beat', 'start_script', 'other'])
            else:
                log(f"无效选项: {choice}", "WARNING")

        # 去重
        selected_categories = list(set(selected_categories))
        return selected_categories

    except Exception as e:
        log(f"解析选择时出错: {e}", "ERROR")
        return []

def main():
    print_banner()

    log("正在查找Python进程...", "INFO")
    processes = get_python_processes_detailed()

    if not processes:
        log("没有找到Python进程", "INFO")
        return

    # 按类型分类进程
    categories = categorize_processes(processes)

    # 显示进程分类
    total_count = display_processes_by_category(categories)

    if total_count == 0:
        log("没有找到相关的Python进程", "INFO")
        return

    log(f"\n总共找到 {total_count} 个Python进程", "INFO")
    log("-" * 60, "INFO")

    while True:
        try:
            selection = show_menu()

            if not selection or selection == '0':
                log("操作已取消", "INFO")
                break

            selected_categories = parse_menu_selection(selection, categories)

            if selected_categories is None:  # 用户选择退出
                log("操作已取消", "INFO")
                break

            if not selected_categories:
                log("没有有效的选择，请重新选择", "WARNING")
                continue

            # 确认操作
            category_names = {
                'fastapi': 'FastAPI 服务',
                'celery_worker': 'Celery Worker 进程',
                'celery_beat': 'Celery Beat 调度器',
                'start_script': '启动脚本',
                'other': '其他Python进程'
            }

            services_to_stop = [category_names.get(cat, cat) for cat in selected_categories]
            log(f"\n准备停止以下服务: {', '.join(services_to_stop)}", "WARNING")

            confirm = input("确认停止? (y/N): ").strip().lower()
            if confirm not in ['y', 'yes', '是']:
                log("操作已取消", "INFO")
                continue

            # 执行停止操作
            success = stop_processes_by_category(categories, selected_categories)

            if success:
                # 验证进程是否已停止
                time.sleep(2)
                log("\n验证停止结果...", "INFO")
                remaining_processes = get_python_processes_detailed()
                remaining_categories = categorize_processes(remaining_processes)

                still_running = []
                for cat in selected_categories:
                    if remaining_categories.get(cat):
                        still_running.append(category_names.get(cat, cat))

                if still_running:
                    log(f"以下服务仍在运行: {', '.join(still_running)}", "WARNING")
                else:
                    log("所有选定的服务已成功停止", "SUCCESS")

            # 询问是否继续
            continue_choice = input("\n是否继续停止其他服务? (y/N): ").strip().lower()
            if continue_choice not in ['y', 'yes', '是']:
                break

        except KeyboardInterrupt:
            log("\n操作被用户中断", "INFO")
            break
        except Exception as e:
            log(f"操作过程中出现错误: {e}", "ERROR")
            break

    log("脚本执行完成", "SUCCESS")

if __name__ == '__main__':
    main()
